#ifndef ROADREPAIRER_H
#define ROADREPAIRER_H

#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>
#include <string>
#include <vector>
#include <cmath>

class RoadRepairer
{
public:
    RoadRepairer();
    void initialize(const std::string& modelPath, bool useGPU = false);
    cv::Mat repair(const cv::Mat& image, const cv::Mat& coarseMask);

private:
    cv::Mat preprocess_image(const cv::Mat& image);
    cv::Mat preprocess_mask(const cv::Mat& mask);
    cv::Mat postprocess(const float* outputData, const std::vector<int64_t>& outputShape);
    std::vector<float> prepare_input_tensor(const cv::Mat& processedImage, const cv::Mat& processedMask);

    static constexpr int INPUT_HEIGHT = 512;
    static constexpr int INPUT_WIDTH = 512;
    static constexpr int CHANNELS = 3;
    const std::vector<float> MEAN = { 123.675f, 116.28f, 103.53f };
    const std::vector<float> STD = { 58.395f, 57.12f, 57.375f };

    Ort::Env env{ nullptr };
    Ort::SessionOptions sessionOptions{ nullptr };
    Ort::Session session{ nullptr };
    std::vector<const char*> inputNames;
    std::vector<const char*> outputNames;
    std::vector<int64_t> InputShape;
    std::vector<int64_t> OutputShape;
};

#endif